// Copyright (C) 2016 Google Inc.
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//
////////////////////////////////////////////////////////////////////////////////

#ifndef DML_DEEPMIND_LUA_READ_H_
#define DML_DEEPMIND_LUA_READ_H_

#include <array>
#include <cstdlib>
#include <sstream>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>

#include "deepmind/lua/lua.h"

namespace deepmind {
namespace lab {
namespace lua {

template <typename Type>
Type* ReadUDT(lua_State* L, int idx, const char* tname) {
  if (!lua_isuserdata(L, idx)) {
    return nullptr;
  }
  void* p = lua_touserdata(L, idx);
  Type* t = nullptr;

  if (lua_getmetatable(L, idx)) {
    lua_pushstring(L, tname);
    lua_gettable(L, LUA_REGISTRYINDEX);
    if (lua_rawequal(L, -1, -2)) {
      t = static_cast<Type*>(p);
    }
    lua_pop(L, 2);
  }

  return t;
}

// In all Read overloads, '*result' is filled in if value at the stack
// location 'idx' is valid for the given type. The return value indicates
// whether the read was valid. If the read fails, '*result' is unmodified.

inline bool Read(lua_State* L, int idx, std::string* result) {
  if (lua_type(L, idx) == LUA_TSTRING) {
    std::size_t length = 0;
    const char* result_cstr = lua_tolstring(L, idx, &length);
    *result = std::string(result_cstr, length);
    return true;
  }
  return false;
}

inline bool Read(lua_State* L, int idx, bool* result) {
  if (lua_type(L, idx) == LUA_TBOOLEAN) {
    *result = lua_toboolean(L, idx);
    return true;
  }
  return false;
}

inline bool Read(lua_State* L, int idx, lua_Integer* result) {
  if (lua_type(L, idx) == LUA_TNUMBER) {
    *result = lua_tointeger(L, idx);
    return true;
  }
  return false;
}

inline bool Read(lua_State* L, int idx, lua_Number* result) {
  if (lua_type(L, idx) == LUA_TNUMBER) {
    *result = lua_tonumber(L, idx);
    return true;
  }
  return false;
}

// Convenience wrapper for arbitrary signed integral types.
template <typename T>
typename std::enable_if<std::is_integral<T>::value &&
                        std::is_signed<T>::value, bool>::type
Read(lua_State* L, int idx, T* out) {
  lua_Integer result;
  if (!Read(L, idx, &result)) {
    return false;
  } else {
    *out = result;
    return true;
  }
}

// Convenience wrapper for arbitrary unsigned integral types.
// Read fails if value is negative.
template <typename T>
typename std::enable_if<std::is_unsigned<T>::value &&
                        !std::is_same<T, bool>::value, bool>::type
Read(lua_State* L, int idx, T* out) {
  lua_Integer result;
  if (!Read(L, idx, &result) || result < 0) {
    return false;
  } else {
    *out = result;
    return true;
  }
}

// Convenience wrapper for arbitrary floating-point types.
template <typename T>
typename std::enable_if<std::is_floating_point<T>::value, bool>::type
Read(lua_State* L, int idx, T* out) {
  lua_Number result;
  if (!Read(L, idx, &result)) {
    return false;
  } else {
    *out = result;
    return true;
  }
}

inline bool Read(lua_State* L, int idx, lua_CFunction* result) {
  if (lua_type(L, idx) == LUA_TFUNCTION) {
    *result = lua_tocfunction(L, idx);
    return true;
  }
  return false;
}

template <typename T>
bool Read(lua_State* L, int idx, T** result) {
  if (lua_type(L, idx) == LUA_TLIGHTUSERDATA) {
    *result = static_cast<T*>(lua_touserdata(L, idx));
    return true;
  }
  return false;
}

// Reads an array from the Lua stack. On success, the array is stored in
// '*result'; on failure, '*result' is unmodified. Returns whether the function
// succeeds.
//
// The function fails if the value on the stack at position 'idx' is not a
// table, or if any array element of the table is not readable as type T.
// Non-array table elements are silently ignored.
template <typename T, typename A>
bool Read(lua_State* L, int idx, std::vector<T, A>* result);

// Reads a Lua array into '*result'. The failure conditions are the same as in
// the previous function, but '*result' may be modified even if this function
// fails.
template <typename T, std::size_t N>
bool Read(lua_State* L, int idx, std::array<T, N>* values);

// Reads a table from the Lua stack. On success, the table is stored in
// '*result'; on failure, '*result' is unmodified. Returns whether the function
// succeeds.
//
// The function fails if the value on the stack at position 'idx' is not a
// table, or if the table contains an entry whose key cannot be read as K or
// whose value cannot be read as T.
template <typename K, typename T, typename H, typename C, typename A>
bool Read(lua_State* L, int idx, std::unordered_map<K, T, H, C, A>* result);

template <typename T, std::size_t N>
bool Read(lua_State* L, int idx, std::array<T, N>* values) {
  if (lua_type(L, idx) == LUA_TTABLE) {
    std::size_t count = ArrayLength(L, idx);
    if (count >= values->size()) {
      for (std::size_t i = 0; i < values->size(); ++i) {
        lua_rawgeti(L, idx, i + 1);
        if (!Read(L, -1, &(*values)[i])) {
          lua_pop(L, 1);
          return false;
        }
        lua_pop(L, 1);
      }
      return true;
    }
  }
  return false;
}

template <typename T, typename A>
bool Read(lua_State* L, int idx, std::vector<T, A>* result) {
  std::vector<T, A> local_result;
  if (lua_type(L, idx) == LUA_TTABLE) {
    std::size_t count = ArrayLength(L, idx);
    local_result.reserve(count);
    for (std::size_t i = 0; i < count; ++i) {
      lua_rawgeti(L, idx, i + 1);
      T value;
      if (Read(L, -1, &value)) {
        local_result.push_back(std::move(value));
      } else {
        lua_pop(L, 1);
        return false;
      }
      lua_pop(L, 1);
    }
    result->swap(local_result);
    return true;
  }
  return false;
}

template <typename K, typename T, typename H, typename C, typename A>
bool Read(lua_State* L, int idx, std::unordered_map<K, T, H, C, A>* result) {
  std::unordered_map<K, T, H, C, A> local_result;
  if (lua_type(L, idx) == LUA_TTABLE) {
    if (idx < 0) {
      idx = lua_gettop(L) + idx + 1;
    }
    lua_pushnil(L);
    while (lua_next(L, idx) != 0) {
      K key;
      if (!Read(L, -2, &key)) {
        lua_pop(L, 2);
        return false;
      }
      T value;
      if (!Read(L, -1, &value)) {
        lua_pop(L, 2);
        return false;
      }
      local_result.emplace(std::move(key), std::move(value));
      lua_pop(L, 1);
    }
    result->swap(local_result);
    return true;
  }
  return false;
}

// Coerce result into human readable string. (Does not fail.)
inline std::string ToString(lua_State* L, int idx) {
  std::stringstream ss;
  switch (lua_type(L, idx)) {
    case LUA_TNONE:
      ss << "(none)";
      break;
    case LUA_TNIL:
      ss << "(nil)";
      break;
    case LUA_TBOOLEAN:
      ss << (lua_toboolean(L, idx) ? "true" : "false");
      break;
    case LUA_TLIGHTUSERDATA:
      ss << "pointer [" << lua_touserdata(L, idx) << "]";
      break;
    case LUA_TNUMBER:
      ss << lua_tonumber(L, idx);
      break;
    case LUA_TSTRING: {
      std::size_t length = 0;
      const char* result_cstr = lua_tolstring(L, idx, &length);
      ss << std::string(result_cstr, length);
      break;
    }
    case LUA_TTABLE:
      ss << "(table)";
      break;
    case LUA_TFUNCTION:
      ss << "function [" << lua_tocfunction(L, idx) << "]";
      break;
    case LUA_TUSERDATA:
      ss << "user pointer [" << lua_touserdata(L, idx) << "]";
      break;
    default:
      ss << "(unknown)";
      break;
  }
  return ss.str();
}

}  // namespace lua
}  // namespace lab
}  // namespace deepmind

#endif  // DML_DEEPMIND_LUA_READ_H_
